import gradio as gr
import cv2 as cv
import matplotlib.pyplot as plt
import numpy as np
from pinecone import Pinecone
import logging

# 配置日志
# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')

##### 1. 初始化Pinecone #####
# Initialize Pinecone

# 替换为你的Pinecone API密钥
# Replace with your Pinecone API key
pc = Pinecone(api_key="67ad532a-d0c4-4aca-9b00-a7ea204cc919")

# 连接到你的索引
# Connect to your index
index_name = "mnist-index"
index = pc.Index(index_name)

logging.info(f"已成功连接到索引 '{index_name}'。")
# Successfully connected to index '{index_name}'.

##### 2. 定义预测函数 #####
# Define prediction function

def predict(image):
    # 检查图像是否为空
    # Check if the image is empty
    if image is None:
        return None
    
    # 转换图像为灰度图
    # Convert image to grayscale
    image = cv.cvtColor(image["composite"], cv.COLOR_RGBA2GRAY)
    
    # 保存图像用于调试
    # Save image for debugging
    plt.imsave("image1.png", image, cmap='gray')

    # 调整图像大小为28x28并展平
    # Resize to 28x28 and flatten the image
    image = cv.resize(image, (28, 28), interpolation=cv.INTER_AREA)
    plt.imsave("image2.png", image, cmap="gray")
    
    # 归一化并添加epsilon
    # Normalize and add epsilon
    epsilon = 1e-8
    image = image.astype('float32') / 255.0 + epsilon
    
    image = image.ravel()

    # 查询Pinecone
    # Query Pinecone
    max_retries = 3
    for attempt in range(max_retries):
        try:
            query_result = index.query(
                vector=image.tolist(),
                top_k=11,  # 使用k=11，与原始代码一致
                include_metadata=True
            )
            break
        except Exception as e:
            if attempt < max_retries - 1:
                logging.warning(f"查询失败，正在重试（第{attempt+1}次）: {str(e)}")
                # Query failed, retrying (attempt {attempt+1}): {str(e)}
            else:
                logging.error(f"查询失败，已达到最大重试次数: {str(e)}")
                # Query failed, maximum retries reached: {str(e)}
                return "查询失败"  # Query failed

    # 从查询结果中提取预测标签
    # Extract predicted labels from query result
    nearest_labels = [match['metadata']['label'] for match in query_result['matches']]
    
    # 选择出现次数最多的标签作为预测结果
    # Choose the most common label as the prediction
    predicted_label = max(set(nearest_labels), key=nearest_labels.count)

    return str(predicted_label)

##### 3. 构建Gradio界面 #####
# Build Gradio interface

gr.Interface(fn=predict, inputs=gr.ImageEditor(), outputs='label').launch()